package org.dromara.llm.llm.factory;

import lombok.extern.slf4j.Slf4j;
import okhttp3.*;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.dromara.llm.domain.SysLlmConfig;
import org.springframework.stereotype.Component;
import org.springframework.web.multipart.MultipartFile;

import java.io.IOException;
import java.util.Base64;

@Component
@Slf4j
public class LLMImageOcr {
    // 大模型 API 配置（根据实际模型替换）
    private String API_URL = ""; // 大模型的API接口地址,例如：https://open.bigmodel.cn/api/paas/v4
    private String API_KEY = ""; // 你的API密钥,替换为实际密钥
    private String MODEL = ""; // 模型名称, 例如：gpt-4-vision-preview 或 qwen-vl-ocr-latest 或 GLM-4V-Plus

    public void initLLMImage(SysLlmConfig config) {
        // 直接获取模型名称，无需转换
        MODEL = config.getConfigName();
        // 获取API的URL
        API_URL = config.getApiUrl() + "/chat/completions";
        // 获取API密钥
        API_KEY = config.getApiKey();
        log.info("初始化LLMImage模型：{}" ,config);
    }

    /**
     * 识别图片中的文字
     * @param imageUrl 图片URL
     * @param prompt 提示词 如："识别图片中的所有文字，直接输出结果，不要额外说明"
     * @return 识别结果
     * @throws IOException 异常
     */
    public String recognizeImageText(String imageUrl,String prompt) throws IOException {
        // 1. 构造请求体
        ObjectMapper mapper = new ObjectMapper();
        // 构建提示词和图片信息
        JsonNode requestBody = mapper.createObjectNode()
            .put("model", MODEL)
            .put("max_tokens", 2000) // 确保足够容纳识别结果
            .set("messages", mapper.createArrayNode()
                .add(mapper.createObjectNode()
                    .put("role", "user")
                    .set("content", mapper.createArrayNode()
                        // 提示词
                        .add(mapper.createObjectNode()
                            .put("type", "text")
                            .put("text", prompt))
                        // 图片信息（URL格式）
                        .add(mapper.createObjectNode()
                            .put("type", "image_url")
                            // 如果是千问的模型，如果为url，直接使用image_url指向图片的url地址，base64如下需要加一个url节点
                            //url方式："image_url": "https://img.alicdn.com/imgextra/i2/O1CN01ktT8451iQutqReELT_!!6000000004408-0-tps-689-487.jpg"
                            // 智普AI的base64方式和url方式都是一样的，结构固定
                            .set("image_url", mapper.createObjectNode()
                                .put("url", imageUrl))
                        )
                    )));

        // 2. 发送请求
        OkHttpClient client = new OkHttpClient();
        Request request = new Request.Builder()
            .url(API_URL)
            .addHeader("Content-Type", "application/json")
            .addHeader("Authorization", "Bearer " + API_KEY) // 认证头（不同模型可能有差异）
            .post(RequestBody.create(
                MediaType.parse("application/json"),
                mapper.writeValueAsString(requestBody)
                        ))
            .build();

        // 3. 解析响应
        try (Response response = client.newCall(request).execute()) {
            if (!response.isSuccessful() || response.body() == null) {
                throw new IOException("请求失败：" + response.code());
            }
            JsonNode responseNode = mapper.readTree(response.body().string());
            // 提取文字（不同模型的响应路径不同，需根据实际调整）
            return responseNode.get("choices").get(0).get("message").get("content").asText();
        }
    }

    /**
     * 识别图片中的文字
     * @param file 图片文件，转换为base64编码
     * @param prompt 提示词 如："识别图片中的所有文字，直接输出结果，不要额外说明"
     * @return 识别结果
     * @throws IOException 异常
     */
    public String recognizeImageText(MultipartFile file, String prompt) throws IOException {
        String contentType = file.getContentType();
        String preBase64Image ="";
        if (contentType != null) {
            if (contentType.startsWith("image/")) {
                preBase64Image = "data:"+contentType+";base64,";
            } else {
                throw new IOException("不支持的文件格式");
            }
        }
        String base64Image = Base64.getEncoder().encodeToString(file.getInputStream().readAllBytes());
        // 需要注意，传入Base64，图像格式（即image/{format}）需要与支持的图片列表中的Content Type保持一致。
        // PNG图像：  “data:image/png;base64,{base64_image}"
        // JPEG图像： "data:image/jpeg;base64,{base64_image}"
        // WEBP图像： "data:image/webp;base64,{base64_image}"
        // image/tiff image/heic image/bmp
        base64Image = preBase64Image + base64Image;
        return recognizeImageText(base64Image, prompt);
    }

}
